{
 "cells": [
  {
   "cell_type": "raw",
   "id": "b61d713f-3ef8-4bef-8ff1-9859276c7726",
   "metadata": {},
   "source": [
    "[2024-09-26]\n",
    "pretrain 데이터 전처리 실행"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0a969593-842a-469d-8ce8-c61518679c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import wfdb\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "_LEAD_NAMES = [\"I\", \"II\", \"III\", \"aVR\", \"aVL\", \"aVF\", \"V1\", \"V2\", \"V3\", \"V4\", \"V5\", \"V6\"]\n",
    "# _LEAD_NAMES = ['I', 'II', 'III', 'AVR', 'AVL', 'AVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6']"
   ]
  },
  {
   "cell_type": "raw",
   "id": "fb898884-9912-4e53-b6f7-564962cbe1d8",
   "metadata": {},
   "source": [
    "if __name__ == \"__main__\":\n",
    "    run(get_parser())\n",
    "\n",
    "python process_ecg.py \\\n",
    "    --input_dir ${DATABASE_DIRECTORY} \\\n",
    "    --output_dir ${WAVEFORM_DIRECTORY} \\\n",
    "    --index_path ${INDEX_PATH}\n",
    "\n",
    "- input_dir: 원시 WFDB 파형 파일이 저장된 directory\n",
    "- output_dir: ECG PKL 파일 저장될 directory\n",
    "- index_path: index csv 파일 저장될 경로"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "319dfb1f-5ac8-46d3-8d3a-a9038939a83b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_parser():\n",
    "    description = \"Process WFDB ECG database.\"\n",
    "    # parser = argparse.ArgumentParser(description=description)\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument('-i',\n",
    "                        '--input_dir',\n",
    "                        type=str,\n",
    "                        # required=True,\n",
    "                        default='/tf/physionet.org/files/challenge-2021/1.0.3/training/chapman_shaoxing/',\n",
    "                        help=\"Path to the WFDB ECG database directory.\")\n",
    "    parser.add_argument('-o',\n",
    "                        '--output_dir',\n",
    "                        type=str,\n",
    "                        # required=True,\n",
    "                        default='./chapman/ecgs/',\n",
    "                        help=\"Path to the directory where the preprocessed signals will be saved.\")\n",
    "    parser.add_argument('--index_path',\n",
    "                        type=str,\n",
    "                        default='./chapman/index.csv',\n",
    "                        help=\"Path to the index file.\")\n",
    "    args = parser.parse_args(\"\")\n",
    "    return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "41ae9fa0-ad53-4fa9-a5c6-e9163bddd7a2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Namespace(input_dir='/tf/physionet.org/files/challenge-2021/1.0.3/training/chapman_shaoxing/', output_dir='./chapman/ecgs/', index_path='./chapman/index.csv')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "args = get_parser()\n",
    "args"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53e06139-70d8-4e27-a8d7-3ccf531599ee",
   "metadata": {},
   "source": [
    "get_parser() 부분 실행 완 => args 처리됨. 이후 run 코드 순서대로 실행하면 됨!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ea547cef-aab4-487c-8911-56393c721ad4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_records(root_dir):\n",
    "    \"\"\"Find all the .hea files in the root directory and its subdirectories.\n",
    "    Args:\n",
    "        root_dir (str): The directory to search for .hea files.\n",
    "    Returns:\n",
    "        records (set): A set of record names.\n",
    "                       (e.g., ['database/1/ecg001', 'database/1/ecg001', ..., 'database/9/ecg991'])\n",
    "    \"\"\"\n",
    "    records = set()\n",
    "    for root, _, files in os.walk(root_dir):\n",
    "        for file in files:\n",
    "            extension = os.path.splitext(file)[1]\n",
    "            if extension == '.hea':\n",
    "                record = os.path.relpath(os.path.join(root, file), root_dir)[:-4]\n",
    "                records.add(record)\n",
    "    records = sorted(records)\n",
    "    return records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a348247-59b8-4609-b0b9-c19bf09fd3f3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['g1/JS00001',\n",
       " 'g1/JS00002',\n",
       " 'g1/JS00004',\n",
       " 'g1/JS00005',\n",
       " 'g1/JS00006',\n",
       " 'g1/JS00007',\n",
       " 'g1/JS00008',\n",
       " 'g1/JS00009',\n",
       " 'g1/JS00010',\n",
       " 'g1/JS00011',\n",
       " 'g1/JS00012',\n",
       " 'g1/JS00013',\n",
       " 'g1/JS00014',\n",
       " 'g1/JS00015',\n",
       " 'g1/JS00016',\n",
       " 'g1/JS00017',\n",
       " 'g1/JS00018',\n",
       " 'g1/JS00019',\n",
       " 'g1/JS00020',\n",
       " 'g1/JS00021',\n",
       " 'g1/JS00022',\n",
       " 'g1/JS00023',\n",
       " 'g1/JS00024',\n",
       " 'g1/JS00025',\n",
       " 'g1/JS00026',\n",
       " 'g1/JS00027',\n",
       " 'g1/JS00029',\n",
       " 'g1/JS00030',\n",
       " 'g1/JS00031',\n",
       " 'g1/JS00032',\n",
       " 'g1/JS00033',\n",
       " 'g1/JS00034',\n",
       " 'g1/JS00036',\n",
       " 'g1/JS00037',\n",
       " 'g1/JS00038',\n",
       " 'g1/JS00039',\n",
       " 'g1/JS00040',\n",
       " 'g1/JS00041',\n",
       " 'g1/JS00042',\n",
       " 'g1/JS00043',\n",
       " 'g1/JS00044',\n",
       " 'g1/JS00045',\n",
       " 'g1/JS00046',\n",
       " 'g1/JS00047',\n",
       " 'g1/JS00048',\n",
       " 'g1/JS00049',\n",
       " 'g1/JS00050',\n",
       " 'g1/JS00051',\n",
       " 'g1/JS00052',\n",
       " 'g1/JS00053',\n",
       " 'g1/JS00054',\n",
       " 'g1/JS00055',\n",
       " 'g1/JS00056',\n",
       " 'g1/JS00057',\n",
       " 'g1/JS00058',\n",
       " 'g1/JS00059',\n",
       " 'g1/JS00060',\n",
       " 'g1/JS00061',\n",
       " 'g1/JS00062',\n",
       " 'g1/JS00063',\n",
       " 'g1/JS00064',\n",
       " 'g1/JS00065',\n",
       " 'g1/JS00066',\n",
       " 'g1/JS00067',\n",
       " 'g1/JS00068',\n",
       " 'g1/JS00069',\n",
       " 'g1/JS00070',\n",
       " 'g1/JS00071',\n",
       " 'g1/JS00072',\n",
       " 'g1/JS00073',\n",
       " 'g1/JS00074',\n",
       " 'g1/JS00075',\n",
       " 'g1/JS00076',\n",
       " 'g1/JS00077',\n",
       " 'g1/JS00078',\n",
       " 'g1/JS00079',\n",
       " 'g1/JS00080',\n",
       " 'g1/JS00081',\n",
       " 'g1/JS00082',\n",
       " 'g1/JS00083',\n",
       " 'g1/JS00084',\n",
       " 'g1/JS00085',\n",
       " 'g1/JS00086',\n",
       " 'g1/JS00087',\n",
       " 'g1/JS00088',\n",
       " 'g1/JS00089',\n",
       " 'g1/JS00090',\n",
       " 'g1/JS00091',\n",
       " 'g1/JS00092',\n",
       " 'g1/JS00093',\n",
       " 'g1/JS00094',\n",
       " 'g1/JS00095',\n",
       " 'g1/JS00096',\n",
       " 'g1/JS00097',\n",
       " 'g1/JS00099',\n",
       " 'g1/JS00100',\n",
       " 'g1/JS00101',\n",
       " 'g1/JS00102',\n",
       " 'g1/JS00103',\n",
       " 'g1/JS00104',\n",
       " 'g1/JS00105',\n",
       " 'g1/JS00106',\n",
       " 'g1/JS00107',\n",
       " 'g1/JS00108',\n",
       " 'g1/JS00109',\n",
       " 'g1/JS00110',\n",
       " 'g1/JS00111',\n",
       " 'g1/JS00112',\n",
       " 'g1/JS00113',\n",
       " 'g1/JS00114',\n",
       " 'g1/JS00115',\n",
       " 'g1/JS00116',\n",
       " 'g1/JS00117',\n",
       " 'g1/JS00118',\n",
       " 'g1/JS00119',\n",
       " 'g1/JS00120',\n",
       " 'g1/JS00121',\n",
       " 'g1/JS00122',\n",
       " 'g1/JS00123',\n",
       " 'g1/JS00124',\n",
       " 'g1/JS00125',\n",
       " 'g1/JS00126',\n",
       " 'g1/JS00127',\n",
       " 'g1/JS00128',\n",
       " 'g1/JS00129',\n",
       " 'g1/JS00130',\n",
       " 'g1/JS00131',\n",
       " 'g1/JS00132',\n",
       " 'g1/JS00133',\n",
       " 'g1/JS00134',\n",
       " 'g1/JS00137',\n",
       " 'g1/JS00138',\n",
       " 'g1/JS00139',\n",
       " 'g1/JS00140',\n",
       " 'g1/JS00141',\n",
       " 'g1/JS00142',\n",
       " 'g1/JS00143',\n",
       " 'g1/JS00144',\n",
       " 'g1/JS00145',\n",
       " 'g1/JS00146',\n",
       " 'g1/JS00147',\n",
       " 'g1/JS00148',\n",
       " 'g1/JS00150',\n",
       " 'g1/JS00151',\n",
       " 'g1/JS00152',\n",
       " 'g1/JS00153',\n",
       " 'g1/JS00154',\n",
       " 'g1/JS00155',\n",
       " 'g1/JS00156',\n",
       " 'g1/JS00157',\n",
       " 'g1/JS00158',\n",
       " 'g1/JS00159',\n",
       " 'g1/JS00160',\n",
       " 'g1/JS00161',\n",
       " 'g1/JS00162',\n",
       " 'g1/JS00163',\n",
       " 'g1/JS00164',\n",
       " 'g1/JS00165',\n",
       " 'g1/JS00166',\n",
       " 'g1/JS00167',\n",
       " 'g1/JS00168',\n",
       " 'g1/JS00169',\n",
       " 'g1/JS00170',\n",
       " 'g1/JS00171',\n",
       " 'g1/JS00172',\n",
       " 'g1/JS00173',\n",
       " 'g1/JS00174',\n",
       " 'g1/JS00175',\n",
       " 'g1/JS00176',\n",
       " 'g1/JS00177',\n",
       " 'g1/JS00178',\n",
       " 'g1/JS00179',\n",
       " 'g1/JS00180',\n",
       " 'g1/JS00181',\n",
       " 'g1/JS00182',\n",
       " 'g1/JS00183',\n",
       " 'g1/JS00184',\n",
       " 'g1/JS00185',\n",
       " 'g1/JS00186',\n",
       " 'g1/JS00187',\n",
       " 'g1/JS00188',\n",
       " 'g1/JS00189',\n",
       " 'g1/JS00190',\n",
       " 'g1/JS00191',\n",
       " 'g1/JS00192',\n",
       " 'g1/JS00193',\n",
       " 'g1/JS00194',\n",
       " 'g1/JS00195',\n",
       " 'g1/JS00196',\n",
       " 'g1/JS00197',\n",
       " 'g1/JS00198',\n",
       " 'g1/JS00199',\n",
       " 'g1/JS00200',\n",
       " 'g1/JS00201',\n",
       " 'g1/JS00202',\n",
       " 'g1/JS00203',\n",
       " 'g1/JS00204',\n",
       " 'g1/JS00205',\n",
       " 'g1/JS00206',\n",
       " 'g1/JS00207',\n",
       " 'g1/JS00208',\n",
       " 'g1/JS00209',\n",
       " 'g1/JS00210',\n",
       " 'g1/JS00211',\n",
       " 'g1/JS00212',\n",
       " 'g1/JS00213',\n",
       " 'g1/JS00215',\n",
       " 'g1/JS00216',\n",
       " 'g1/JS00217',\n",
       " 'g1/JS00218',\n",
       " 'g1/JS00219',\n",
       " 'g1/JS00220',\n",
       " 'g1/JS00221',\n",
       " 'g1/JS00222',\n",
       " 'g1/JS00223',\n",
       " 'g1/JS00224',\n",
       " 'g1/JS00225',\n",
       " 'g1/JS00226',\n",
       " 'g1/JS00227',\n",
       " 'g1/JS00228',\n",
       " 'g1/JS00229',\n",
       " 'g1/JS00230',\n",
       " 'g1/JS00231',\n",
       " 'g1/JS00232',\n",
       " 'g1/JS00233',\n",
       " 'g1/JS00234',\n",
       " 'g1/JS00235',\n",
       " 'g1/JS00236',\n",
       " 'g1/JS00237',\n",
       " 'g1/JS00238',\n",
       " 'g1/JS00239',\n",
       " 'g1/JS00240',\n",
       " 'g1/JS00241',\n",
       " 'g1/JS00242',\n",
       " 'g1/JS00243',\n",
       " 'g1/JS00244',\n",
       " 'g1/JS00245',\n",
       " 'g1/JS00247',\n",
       " 'g1/JS00248',\n",
       " 'g1/JS00249',\n",
       " 'g1/JS00250',\n",
       " 'g1/JS00251',\n",
       " 'g1/JS00252',\n",
       " 'g1/JS00253',\n",
       " 'g1/JS00254',\n",
       " 'g1/JS00255',\n",
       " 'g1/JS00256',\n",
       " 'g1/JS00257',\n",
       " 'g1/JS00258',\n",
       " 'g1/JS00259',\n",
       " 'g1/JS00260',\n",
       " 'g1/JS00261',\n",
       " 'g1/JS00262',\n",
       " 'g1/JS00263',\n",
       " 'g1/JS00264',\n",
       " 'g1/JS00265',\n",
       " 'g1/JS00266',\n",
       " 'g1/JS00267',\n",
       " 'g1/JS00268',\n",
       " 'g1/JS00269',\n",
       " 'g1/JS00270',\n",
       " 'g1/JS00271',\n",
       " 'g1/JS00272',\n",
       " 'g1/JS00273',\n",
       " 'g1/JS00274',\n",
       " 'g1/JS00276',\n",
       " 'g1/JS00277',\n",
       " 'g1/JS00278',\n",
       " 'g1/JS00279',\n",
       " 'g1/JS00280',\n",
       " 'g1/JS00281',\n",
       " 'g1/JS00282',\n",
       " 'g1/JS00283',\n",
       " 'g1/JS00284',\n",
       " 'g1/JS00285',\n",
       " 'g1/JS00286',\n",
       " 'g1/JS00287',\n",
       " 'g1/JS00288',\n",
       " 'g1/JS00289',\n",
       " 'g1/JS00290',\n",
       " 'g1/JS00291',\n",
       " 'g1/JS00292',\n",
       " 'g1/JS00293',\n",
       " 'g1/JS00294',\n",
       " 'g1/JS00295',\n",
       " 'g1/JS00296',\n",
       " 'g1/JS00297',\n",
       " 'g1/JS00298',\n",
       " 'g1/JS00299',\n",
       " 'g1/JS00300',\n",
       " 'g1/JS00301',\n",
       " 'g1/JS00302',\n",
       " 'g1/JS00303',\n",
       " 'g1/JS00304',\n",
       " 'g1/JS00305',\n",
       " 'g1/JS00306',\n",
       " 'g1/JS00307',\n",
       " 'g1/JS00308',\n",
       " 'g1/JS00309',\n",
       " 'g1/JS00310',\n",
       " 'g1/JS00311',\n",
       " 'g1/JS00312',\n",
       " 'g1/JS00313',\n",
       " 'g1/JS00314',\n",
       " 'g1/JS00315',\n",
       " 'g1/JS00316',\n",
       " 'g1/JS00317',\n",
       " 'g1/JS00318',\n",
       " 'g1/JS00319',\n",
       " 'g1/JS00320',\n",
       " 'g1/JS00321',\n",
       " 'g1/JS00322',\n",
       " 'g1/JS00323',\n",
       " 'g1/JS00324',\n",
       " 'g1/JS00325',\n",
       " 'g1/JS00326',\n",
       " 'g1/JS00327',\n",
       " 'g1/JS00328',\n",
       " 'g1/JS00329',\n",
       " 'g1/JS00330',\n",
       " 'g1/JS00331',\n",
       " 'g1/JS00332',\n",
       " 'g1/JS00333',\n",
       " 'g1/JS00334',\n",
       " 'g1/JS00335',\n",
       " 'g1/JS00336',\n",
       " 'g1/JS00337',\n",
       " 'g1/JS00338',\n",
       " 'g1/JS00339',\n",
       " 'g1/JS00340',\n",
       " 'g1/JS00341',\n",
       " 'g1/JS00342',\n",
       " 'g1/JS00343',\n",
       " 'g1/JS00344',\n",
       " 'g1/JS00345',\n",
       " 'g1/JS00346',\n",
       " 'g1/JS00347',\n",
       " 'g1/JS00348',\n",
       " 'g1/JS00349',\n",
       " 'g1/JS00350',\n",
       " 'g1/JS00351',\n",
       " 'g1/JS00352',\n",
       " 'g1/JS00353',\n",
       " 'g1/JS00354',\n",
       " 'g1/JS00355',\n",
       " 'g1/JS00356',\n",
       " 'g1/JS00357',\n",
       " 'g1/JS00358',\n",
       " 'g1/JS00359',\n",
       " 'g1/JS00360',\n",
       " 'g1/JS00361',\n",
       " 'g1/JS00362',\n",
       " 'g1/JS00363',\n",
       " 'g1/JS00364',\n",
       " 'g1/JS00365',\n",
       " 'g1/JS00366',\n",
       " 'g1/JS00367',\n",
       " 'g1/JS00368',\n",
       " 'g1/JS00369',\n",
       " 'g1/JS00370',\n",
       " 'g1/JS00371',\n",
       " 'g1/JS00372',\n",
       " 'g1/JS00373',\n",
       " 'g1/JS00374',\n",
       " 'g1/JS00375',\n",
       " 'g1/JS00376',\n",
       " 'g1/JS00377',\n",
       " 'g1/JS00378',\n",
       " 'g1/JS00379',\n",
       " 'g1/JS00380',\n",
       " 'g1/JS00381',\n",
       " 'g1/JS00382',\n",
       " 'g1/JS00383',\n",
       " 'g1/JS00384',\n",
       " 'g1/JS00385',\n",
       " 'g1/JS00386',\n",
       " 'g1/JS00387',\n",
       " 'g1/JS00388',\n",
       " 'g1/JS00389',\n",
       " 'g1/JS00390',\n",
       " 'g1/JS00391',\n",
       " 'g1/JS00392',\n",
       " 'g1/JS00393',\n",
       " 'g1/JS00394',\n",
       " 'g1/JS00395',\n",
       " 'g1/JS00396',\n",
       " 'g1/JS00397',\n",
       " 'g1/JS00398',\n",
       " 'g1/JS00399',\n",
       " 'g1/JS00400',\n",
       " 'g1/JS00401',\n",
       " 'g1/JS00402',\n",
       " 'g1/JS00403',\n",
       " 'g1/JS00404',\n",
       " 'g1/JS00405',\n",
       " 'g1/JS00406',\n",
       " 'g1/JS00407',\n",
       " 'g1/JS00408',\n",
       " 'g1/JS00409',\n",
       " 'g1/JS00410',\n",
       " 'g1/JS00411',\n",
       " 'g1/JS00412',\n",
       " 'g1/JS00413',\n",
       " 'g1/JS00414',\n",
       " 'g1/JS00415',\n",
       " 'g1/JS00417',\n",
       " 'g1/JS00420',\n",
       " 'g1/JS00421',\n",
       " 'g1/JS00422',\n",
       " 'g1/JS00423',\n",
       " 'g1/JS00424',\n",
       " 'g1/JS00425',\n",
       " 'g1/JS00426',\n",
       " 'g1/JS00427',\n",
       " 'g1/JS00428',\n",
       " 'g1/JS00429',\n",
       " 'g1/JS00430',\n",
       " 'g1/JS00431',\n",
       " 'g1/JS00432',\n",
       " 'g1/JS00433',\n",
       " 'g1/JS00434',\n",
       " 'g1/JS00435',\n",
       " 'g1/JS00436',\n",
       " 'g1/JS00437',\n",
       " 'g1/JS00438',\n",
       " 'g1/JS00439',\n",
       " 'g1/JS00440',\n",
       " 'g1/JS00441',\n",
       " 'g1/JS00442',\n",
       " 'g1/JS00444',\n",
       " 'g1/JS00445',\n",
       " 'g1/JS00446',\n",
       " 'g1/JS00447',\n",
       " 'g1/JS00448',\n",
       " 'g1/JS00449',\n",
       " 'g1/JS00450',\n",
       " 'g1/JS00451',\n",
       " 'g1/JS00452',\n",
       " 'g1/JS00453',\n",
       " 'g1/JS00454',\n",
       " 'g1/JS00455',\n",
       " 'g1/JS00456',\n",
       " 'g1/JS00457',\n",
       " 'g1/JS00458',\n",
       " 'g1/JS00459',\n",
       " 'g1/JS00461',\n",
       " 'g1/JS00462',\n",
       " 'g1/JS00463',\n",
       " 'g1/JS00464',\n",
       " 'g1/JS00465',\n",
       " 'g1/JS00466',\n",
       " 'g1/JS00467',\n",
       " 'g1/JS00468',\n",
       " 'g1/JS00469',\n",
       " 'g1/JS00470',\n",
       " 'g1/JS00471',\n",
       " 'g1/JS00472',\n",
       " 'g1/JS00473',\n",
       " 'g1/JS00474',\n",
       " 'g1/JS00475',\n",
       " 'g1/JS00476',\n",
       " 'g1/JS00477',\n",
       " 'g1/JS00478',\n",
       " 'g1/JS00479',\n",
       " 'g1/JS00480',\n",
       " 'g1/JS00481',\n",
       " 'g1/JS00482',\n",
       " 'g1/JS00483',\n",
       " 'g1/JS00484',\n",
       " 'g1/JS00485',\n",
       " 'g1/JS00486',\n",
       " 'g1/JS00487',\n",
       " 'g1/JS00488',\n",
       " 'g1/JS00490',\n",
       " 'g1/JS00491',\n",
       " 'g1/JS00492',\n",
       " 'g1/JS00493',\n",
       " 'g1/JS00494',\n",
       " 'g1/JS00495',\n",
       " 'g1/JS00496',\n",
       " 'g1/JS00497',\n",
       " 'g1/JS00498',\n",
       " 'g1/JS00499',\n",
       " 'g1/JS00500',\n",
       " 'g1/JS00501',\n",
       " 'g1/JS00502',\n",
       " 'g1/JS00503',\n",
       " 'g1/JS00504',\n",
       " 'g1/JS00506',\n",
       " 'g1/JS00507',\n",
       " 'g1/JS00508',\n",
       " 'g1/JS00509',\n",
       " 'g1/JS00510',\n",
       " 'g1/JS00511',\n",
       " 'g1/JS00512',\n",
       " 'g1/JS00513',\n",
       " 'g1/JS00514',\n",
       " 'g1/JS00515',\n",
       " 'g1/JS00516',\n",
       " 'g1/JS00517',\n",
       " 'g1/JS00518',\n",
       " 'g1/JS00519',\n",
       " 'g1/JS00520',\n",
       " 'g1/JS00521',\n",
       " 'g1/JS00522',\n",
       " 'g1/JS00523',\n",
       " 'g1/JS00524',\n",
       " 'g1/JS00525',\n",
       " 'g1/JS00526',\n",
       " 'g1/JS00527',\n",
       " 'g1/JS00528',\n",
       " 'g1/JS00529',\n",
       " 'g1/JS00530',\n",
       " 'g1/JS00531',\n",
       " 'g1/JS00532',\n",
       " 'g1/JS00533',\n",
       " 'g1/JS00534',\n",
       " 'g1/JS00535',\n",
       " 'g1/JS00536',\n",
       " 'g1/JS00537',\n",
       " 'g1/JS00538',\n",
       " 'g1/JS00540',\n",
       " 'g1/JS00541',\n",
       " 'g1/JS00543',\n",
       " 'g1/JS00544',\n",
       " 'g1/JS00545',\n",
       " 'g1/JS00546',\n",
       " 'g1/JS00547',\n",
       " 'g1/JS00548',\n",
       " 'g1/JS00549',\n",
       " 'g1/JS00550',\n",
       " 'g1/JS00552',\n",
       " 'g1/JS00553',\n",
       " 'g1/JS00554',\n",
       " 'g1/JS00556',\n",
       " 'g1/JS00557',\n",
       " 'g1/JS00558',\n",
       " 'g1/JS00559',\n",
       " 'g1/JS00560',\n",
       " 'g1/JS00561',\n",
       " 'g1/JS00562',\n",
       " 'g1/JS00563',\n",
       " 'g1/JS00564',\n",
       " 'g1/JS00565',\n",
       " 'g1/JS00566',\n",
       " 'g1/JS00567',\n",
       " 'g1/JS00568',\n",
       " 'g1/JS00569',\n",
       " 'g1/JS00570',\n",
       " 'g1/JS00571',\n",
       " 'g1/JS00572',\n",
       " 'g1/JS00573',\n",
       " 'g1/JS00574',\n",
       " 'g1/JS00575',\n",
       " 'g1/JS00576',\n",
       " 'g1/JS00577',\n",
       " 'g1/JS00578',\n",
       " 'g1/JS00579',\n",
       " 'g1/JS00580',\n",
       " 'g1/JS00581',\n",
       " 'g1/JS00582',\n",
       " 'g1/JS00583',\n",
       " 'g1/JS00585',\n",
       " 'g1/JS00586',\n",
       " 'g1/JS00587',\n",
       " 'g1/JS00588',\n",
       " 'g1/JS00589',\n",
       " 'g1/JS00590',\n",
       " 'g1/JS00591',\n",
       " 'g1/JS00592',\n",
       " 'g1/JS00594',\n",
       " 'g1/JS00595',\n",
       " 'g1/JS00596',\n",
       " 'g1/JS00597',\n",
       " 'g1/JS00598',\n",
       " 'g1/JS00599',\n",
       " 'g1/JS00600',\n",
       " 'g1/JS00601',\n",
       " 'g1/JS00602',\n",
       " 'g1/JS00603',\n",
       " 'g1/JS00605',\n",
       " 'g1/JS00606',\n",
       " 'g1/JS00607',\n",
       " 'g1/JS00608',\n",
       " 'g1/JS00609',\n",
       " 'g1/JS00610',\n",
       " 'g1/JS00611',\n",
       " 'g1/JS00612',\n",
       " 'g1/JS00613',\n",
       " 'g1/JS00614',\n",
       " 'g1/JS00615',\n",
       " 'g1/JS00616',\n",
       " 'g1/JS00617',\n",
       " 'g1/JS00618',\n",
       " 'g1/JS00619',\n",
       " 'g1/JS00620',\n",
       " 'g1/JS00621',\n",
       " 'g1/JS00623',\n",
       " 'g1/JS00624',\n",
       " 'g1/JS00625',\n",
       " 'g1/JS00626',\n",
       " 'g1/JS00627',\n",
       " 'g1/JS00628',\n",
       " 'g1/JS00629',\n",
       " 'g1/JS00630',\n",
       " 'g1/JS00631',\n",
       " 'g1/JS00633',\n",
       " 'g1/JS00634',\n",
       " 'g1/JS00635',\n",
       " 'g1/JS00636',\n",
       " 'g1/JS00637',\n",
       " 'g1/JS00638',\n",
       " 'g1/JS00639',\n",
       " 'g1/JS00640',\n",
       " 'g1/JS00641',\n",
       " 'g1/JS00642',\n",
       " 'g1/JS00643',\n",
       " 'g1/JS00644',\n",
       " 'g1/JS00645',\n",
       " 'g1/JS00646',\n",
       " 'g1/JS00647',\n",
       " 'g1/JS00648',\n",
       " 'g1/JS00649',\n",
       " 'g1/JS00651',\n",
       " 'g1/JS00652',\n",
       " 'g1/JS00653',\n",
       " 'g1/JS00654',\n",
       " 'g1/JS00655',\n",
       " 'g1/JS00656',\n",
       " 'g1/JS00657',\n",
       " 'g1/JS00659',\n",
       " 'g1/JS00660',\n",
       " 'g1/JS00661',\n",
       " 'g1/JS00662',\n",
       " 'g1/JS00663',\n",
       " 'g1/JS00664',\n",
       " 'g1/JS00665',\n",
       " 'g1/JS00666',\n",
       " 'g1/JS00667',\n",
       " 'g1/JS00668',\n",
       " 'g1/JS00669',\n",
       " 'g1/JS00670',\n",
       " 'g1/JS00671',\n",
       " 'g1/JS00672',\n",
       " 'g1/JS00673',\n",
       " 'g1/JS00674',\n",
       " 'g1/JS00675',\n",
       " 'g1/JS00677',\n",
       " 'g1/JS00679',\n",
       " 'g1/JS00680',\n",
       " 'g1/JS00681',\n",
       " 'g1/JS00682',\n",
       " 'g1/JS00683',\n",
       " 'g1/JS00684',\n",
       " 'g1/JS00685',\n",
       " 'g1/JS00686',\n",
       " 'g1/JS00687',\n",
       " 'g1/JS00688',\n",
       " 'g1/JS00689',\n",
       " 'g1/JS00690',\n",
       " 'g1/JS00691',\n",
       " 'g1/JS00692',\n",
       " 'g1/JS00693',\n",
       " 'g1/JS00694',\n",
       " 'g1/JS00695',\n",
       " 'g1/JS00696',\n",
       " 'g1/JS00697',\n",
       " 'g1/JS00698',\n",
       " 'g1/JS00699',\n",
       " 'g1/JS00700',\n",
       " 'g1/JS00701',\n",
       " 'g1/JS00702',\n",
       " 'g1/JS00703',\n",
       " 'g1/JS00704',\n",
       " 'g1/JS00705',\n",
       " 'g1/JS00706',\n",
       " 'g1/JS00707',\n",
       " 'g1/JS00708',\n",
       " 'g1/JS00709',\n",
       " 'g1/JS00711',\n",
       " 'g1/JS00712',\n",
       " 'g1/JS00713',\n",
       " 'g1/JS00714',\n",
       " 'g1/JS00716',\n",
       " 'g1/JS00717',\n",
       " 'g1/JS00718',\n",
       " 'g1/JS00719',\n",
       " 'g1/JS00720',\n",
       " 'g1/JS00721',\n",
       " 'g1/JS00722',\n",
       " 'g1/JS00723',\n",
       " 'g1/JS00724',\n",
       " 'g1/JS00725',\n",
       " 'g1/JS00726',\n",
       " 'g1/JS00727',\n",
       " 'g1/JS00728',\n",
       " 'g1/JS00729',\n",
       " 'g1/JS00731',\n",
       " 'g1/JS00732',\n",
       " 'g1/JS00734',\n",
       " 'g1/JS00735',\n",
       " 'g1/JS00736',\n",
       " 'g1/JS00737',\n",
       " 'g1/JS00738',\n",
       " 'g1/JS00739',\n",
       " 'g1/JS00740',\n",
       " 'g1/JS00741',\n",
       " 'g1/JS00742',\n",
       " 'g1/JS00743',\n",
       " 'g1/JS00744',\n",
       " 'g1/JS00745',\n",
       " 'g1/JS00746',\n",
       " 'g1/JS00747',\n",
       " 'g1/JS00748',\n",
       " 'g1/JS00749',\n",
       " 'g1/JS00750',\n",
       " 'g1/JS00751',\n",
       " 'g1/JS00752',\n",
       " 'g1/JS00753',\n",
       " 'g1/JS00754',\n",
       " 'g1/JS00755',\n",
       " 'g1/JS00756',\n",
       " 'g1/JS00757',\n",
       " 'g1/JS00758',\n",
       " 'g1/JS00759',\n",
       " 'g1/JS00760',\n",
       " 'g1/JS00761',\n",
       " 'g1/JS00762',\n",
       " 'g1/JS00763',\n",
       " 'g1/JS00764',\n",
       " 'g1/JS00765',\n",
       " 'g1/JS00766',\n",
       " 'g1/JS00767',\n",
       " 'g1/JS00768',\n",
       " 'g1/JS00769',\n",
       " 'g1/JS00770',\n",
       " 'g1/JS00771',\n",
       " 'g1/JS00772',\n",
       " 'g1/JS00773',\n",
       " 'g1/JS00774',\n",
       " 'g1/JS00775',\n",
       " 'g1/JS00776',\n",
       " 'g1/JS00777',\n",
       " 'g1/JS00778',\n",
       " 'g1/JS00779',\n",
       " 'g1/JS00780',\n",
       " 'g1/JS00781',\n",
       " 'g1/JS00782',\n",
       " 'g1/JS00783',\n",
       " 'g1/JS00784',\n",
       " 'g1/JS00785',\n",
       " 'g1/JS00786',\n",
       " 'g1/JS00787',\n",
       " 'g1/JS00788',\n",
       " 'g1/JS00790',\n",
       " 'g1/JS00791',\n",
       " 'g1/JS00792',\n",
       " 'g1/JS00793',\n",
       " 'g1/JS00794',\n",
       " 'g1/JS00795',\n",
       " 'g1/JS00796',\n",
       " 'g1/JS00797',\n",
       " 'g1/JS00798',\n",
       " 'g1/JS00799',\n",
       " 'g1/JS00800',\n",
       " 'g1/JS00801',\n",
       " 'g1/JS00803',\n",
       " 'g1/JS00804',\n",
       " 'g1/JS00805',\n",
       " 'g1/JS00806',\n",
       " 'g1/JS00807',\n",
       " 'g1/JS00808',\n",
       " 'g1/JS00809',\n",
       " 'g1/JS00810',\n",
       " 'g1/JS00811',\n",
       " 'g1/JS00812',\n",
       " 'g1/JS00813',\n",
       " 'g1/JS00814',\n",
       " 'g1/JS00815',\n",
       " 'g1/JS00816',\n",
       " 'g1/JS00817',\n",
       " 'g1/JS00818',\n",
       " 'g1/JS00819',\n",
       " 'g1/JS00820',\n",
       " 'g1/JS00821',\n",
       " 'g1/JS00822',\n",
       " 'g1/JS00823',\n",
       " 'g1/JS00824',\n",
       " 'g1/JS00825',\n",
       " 'g1/JS00827',\n",
       " 'g1/JS00828',\n",
       " 'g1/JS00829',\n",
       " 'g1/JS00830',\n",
       " 'g1/JS00831',\n",
       " 'g1/JS00832',\n",
       " 'g1/JS00833',\n",
       " 'g1/JS00834',\n",
       " 'g1/JS00835',\n",
       " 'g1/JS00837',\n",
       " 'g1/JS00838',\n",
       " 'g1/JS00839',\n",
       " 'g1/JS00840',\n",
       " 'g1/JS00841',\n",
       " 'g1/JS00842',\n",
       " 'g1/JS00843',\n",
       " 'g1/JS00844',\n",
       " 'g1/JS00845',\n",
       " 'g1/JS00846',\n",
       " 'g1/JS00847',\n",
       " 'g1/JS00848',\n",
       " 'g1/JS00849',\n",
       " 'g1/JS00850',\n",
       " 'g1/JS00851',\n",
       " 'g1/JS00852',\n",
       " 'g1/JS00853',\n",
       " 'g1/JS00854',\n",
       " 'g1/JS00856',\n",
       " 'g1/JS00857',\n",
       " 'g1/JS00858',\n",
       " 'g1/JS00859',\n",
       " 'g1/JS00860',\n",
       " 'g1/JS00861',\n",
       " 'g1/JS00863',\n",
       " 'g1/JS00864',\n",
       " 'g1/JS00865',\n",
       " 'g1/JS00866',\n",
       " 'g1/JS00867',\n",
       " 'g1/JS00868',\n",
       " 'g1/JS00870',\n",
       " 'g1/JS00871',\n",
       " 'g1/JS00872',\n",
       " 'g1/JS00873',\n",
       " 'g1/JS00875',\n",
       " 'g1/JS00876',\n",
       " 'g1/JS00877',\n",
       " 'g1/JS00878',\n",
       " 'g1/JS00879',\n",
       " 'g1/JS00880',\n",
       " 'g1/JS00881',\n",
       " 'g1/JS00882',\n",
       " 'g1/JS00883',\n",
       " 'g1/JS00884',\n",
       " 'g1/JS00885',\n",
       " 'g1/JS00886',\n",
       " 'g1/JS00887',\n",
       " 'g1/JS00888',\n",
       " 'g1/JS00890',\n",
       " 'g1/JS00891',\n",
       " 'g1/JS00893',\n",
       " 'g1/JS00894',\n",
       " 'g1/JS00895',\n",
       " 'g1/JS00896',\n",
       " 'g1/JS00897',\n",
       " 'g1/JS00898',\n",
       " 'g1/JS00899',\n",
       " 'g1/JS00900',\n",
       " 'g1/JS00901',\n",
       " 'g1/JS00902',\n",
       " 'g1/JS00903',\n",
       " 'g1/JS00904',\n",
       " 'g1/JS00905',\n",
       " 'g1/JS00906',\n",
       " 'g1/JS00907',\n",
       " 'g1/JS00909',\n",
       " 'g1/JS00910',\n",
       " 'g1/JS00911',\n",
       " 'g1/JS00912',\n",
       " 'g1/JS00913',\n",
       " 'g1/JS00914',\n",
       " 'g1/JS00915',\n",
       " 'g1/JS00916',\n",
       " 'g1/JS00917',\n",
       " 'g1/JS00918',\n",
       " 'g1/JS00919',\n",
       " 'g1/JS00921',\n",
       " 'g1/JS00922',\n",
       " 'g1/JS00924',\n",
       " 'g1/JS00925',\n",
       " 'g1/JS00926',\n",
       " 'g1/JS00927',\n",
       " 'g1/JS00928',\n",
       " 'g1/JS00929',\n",
       " 'g1/JS00930',\n",
       " 'g1/JS00931',\n",
       " 'g1/JS00932',\n",
       " 'g1/JS00933',\n",
       " 'g1/JS00934',\n",
       " 'g1/JS00935',\n",
       " 'g1/JS00936',\n",
       " 'g1/JS00937',\n",
       " 'g1/JS00938',\n",
       " 'g1/JS00939',\n",
       " 'g1/JS00940',\n",
       " 'g1/JS00941',\n",
       " 'g1/JS00942',\n",
       " 'g1/JS00943',\n",
       " 'g1/JS00944',\n",
       " 'g1/JS00945',\n",
       " 'g1/JS00946',\n",
       " 'g1/JS00947',\n",
       " 'g1/JS00948',\n",
       " 'g1/JS00949',\n",
       " 'g1/JS00950',\n",
       " 'g1/JS00951',\n",
       " 'g1/JS00952',\n",
       " 'g1/JS00953',\n",
       " 'g1/JS00954',\n",
       " 'g1/JS00955',\n",
       " 'g1/JS00956',\n",
       " 'g1/JS00958',\n",
       " 'g1/JS00959',\n",
       " 'g1/JS00960',\n",
       " 'g1/JS00961',\n",
       " 'g1/JS00962',\n",
       " 'g1/JS00963',\n",
       " 'g1/JS00964',\n",
       " 'g1/JS00965',\n",
       " 'g1/JS00966',\n",
       " 'g1/JS00967',\n",
       " 'g1/JS00968',\n",
       " 'g1/JS00969',\n",
       " 'g1/JS00970',\n",
       " 'g1/JS00971',\n",
       " 'g1/JS00972',\n",
       " 'g1/JS00973',\n",
       " 'g1/JS00974',\n",
       " 'g1/JS00975',\n",
       " 'g1/JS00976',\n",
       " 'g1/JS00977',\n",
       " 'g1/JS00978',\n",
       " 'g1/JS00979',\n",
       " 'g1/JS00980',\n",
       " 'g1/JS00981',\n",
       " 'g1/JS00982',\n",
       " 'g1/JS00983',\n",
       " 'g1/JS00984',\n",
       " 'g1/JS00985',\n",
       " 'g1/JS00986',\n",
       " 'g1/JS00987',\n",
       " 'g1/JS00988',\n",
       " 'g1/JS00989',\n",
       " 'g1/JS00990',\n",
       " 'g1/JS00992',\n",
       " 'g1/JS00993',\n",
       " 'g1/JS00994',\n",
       " 'g1/JS00995',\n",
       " 'g1/JS00996',\n",
       " 'g1/JS00997',\n",
       " 'g1/JS00998',\n",
       " 'g1/JS00999',\n",
       " 'g1/JS01000',\n",
       " 'g1/JS01001',\n",
       " 'g1/JS01002',\n",
       " 'g1/JS01003',\n",
       " 'g1/JS01004',\n",
       " 'g1/JS01005',\n",
       " 'g1/JS01006',\n",
       " 'g1/JS01007',\n",
       " 'g1/JS01008',\n",
       " 'g1/JS01009',\n",
       " 'g1/JS01010',\n",
       " 'g1/JS01011',\n",
       " 'g1/JS01012',\n",
       " 'g1/JS01013',\n",
       " 'g1/JS01014',\n",
       " 'g1/JS01016',\n",
       " 'g1/JS01017',\n",
       " 'g1/JS01018',\n",
       " 'g1/JS01019',\n",
       " 'g1/JS01020',\n",
       " 'g1/JS01021',\n",
       " 'g1/JS01022',\n",
       " 'g1/JS01023',\n",
       " 'g1/JS01024',\n",
       " 'g1/JS01025',\n",
       " 'g1/JS01026',\n",
       " 'g1/JS01028',\n",
       " 'g1/JS01030',\n",
       " 'g1/JS01031',\n",
       " 'g1/JS01032',\n",
       " 'g1/JS01033',\n",
       " 'g1/JS01034',\n",
       " 'g1/JS01035',\n",
       " 'g1/JS01036',\n",
       " 'g1/JS01037',\n",
       " 'g1/JS01038',\n",
       " 'g1/JS01039',\n",
       " 'g1/JS01040',\n",
       " 'g1/JS01041',\n",
       " 'g1/JS01042',\n",
       " 'g1/JS01043',\n",
       " 'g1/JS01044',\n",
       " 'g1/JS01045',\n",
       " 'g1/JS01046',\n",
       " 'g1/JS01047',\n",
       " 'g1/JS01048',\n",
       " 'g1/JS01049',\n",
       " 'g1/JS01050',\n",
       " 'g1/JS01051',\n",
       " 'g10/JS09373',\n",
       " ...]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Identify the header fiels\n",
    "record_rel_paths = find_records(args.input_dir)\n",
    "record_rel_paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f7d2f1b1-edbb-41e5-a649-e567581fe63c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10247,\n",
       " ['g1/JS00001', 'g1/JS00002', 'g1/JS00004', 'g1/JS00005'],\n",
       " ['g9/JS09369', 'g9/JS09370', 'g9/JS09371', 'g9/JS09372'])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(record_rel_paths), record_rel_paths[:4], record_rel_paths[-4:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b5851faa-c9ae-400b-bff6-04a61dd07294",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 10247 records.\n"
     ]
    }
   ],
   "source": [
    "print(f\"Found {len(record_rel_paths)} records.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0d5e3c20-bf34-4f6e-b09e-7b80d868ed7a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>RELATIVE_FILE_PATH</th>\n",
       "      <th>FILE_NAME</th>\n",
       "      <th>SAMPLE_RATE</th>\n",
       "      <th>SOURCE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: [RELATIVE_FILE_PATH, FILE_NAME, SAMPLE_RATE, SOURCE]\n",
       "Index: []"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Prepare an index dataframe\n",
    "index_df = pd.DataFrame(columns = [\"RELATIVE_FILE_PATH\", \"FILE_NAME\", \"SAMPLE_RATE\", \"SOURCE\"])\n",
    "index_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e672e6b6-eafd-4ab2-809e-3b00523dce2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def moving_window_crop(x: np.ndarray, crop_length: int, crop_stride: int) -> np.ndarray:\n",
    "    \"\"\"Crop the input sequence with a moving window.\n",
    "    \"\"\"\n",
    "    if crop_length > x.shape[1]:\n",
    "        raise ValueError(f\"crop_length must be smaller than the length of x ({x.shape[1]}).\")\n",
    "    start_idx = np.arange(0, x.shape[1] - crop_length + 1, crop_stride)\n",
    "    return [x[:, i:i + crop_length] for i in start_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "5f8a34a1-f355-49da-a8b3-c38887eeda2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████| 10247/10247 [48:46<00:00,  3.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved 10247 cropped signals.\n"
     ]
    }
   ],
   "source": [
    "num_saved = 0\n",
    "for record_rel_path in tqdm(record_rel_paths):\n",
    "    record_rel_dir, record_name = os.path.split(record_rel_path)\n",
    "    save_dir = os.path.join(args.output_dir, record_rel_dir)\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    source_name = record_rel_dir.split(\"/\")[0]\n",
    "    signal, record_info = wfdb.rdsamp(os.path.join(args.input_dir, record_rel_path))\n",
    "    lead_idx = np.array([record_info[\"sig_name\"].index(lead_name) for lead_name in _LEAD_NAMES])\n",
    "    signal = signal[:, lead_idx]\n",
    "    fs = record_info[\"fs\"]\n",
    "    signal_length = record_info[\"sig_len\"]\n",
    "    if signal_length < 10 * fs:  # Exclude the ECGs with lengths of less than 10 seconds\n",
    "        continue\n",
    "    cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "    for idx, cropped_signal in enumerate(cropped_signals):\n",
    "        if cropped_signal.shape[1] != 10 * fs or np.isnan(cropped_signal).any():\n",
    "            continue\n",
    "        pd.to_pickle(cropped_signal.astype(np.float32),\n",
    "                     os.path.join(save_dir, f\"{record_name}_{idx}.pkl\"))\n",
    "        index_df.loc[num_saved] = [f\"{record_rel_path}_{idx}.pkl\",\n",
    "                                   f\"{record_name}_{idx}.pkl\",\n",
    "                                   fs,\n",
    "                                   source_name]\n",
    "        num_saved += 1\n",
    "\n",
    "print(f\"Saved {num_saved} cropped signals.\")\n",
    "os.makedirs(os.path.dirname(args.index_path), exist_ok=True)\n",
    "index_df.to_csv(args.index_path, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "bf8b378c-f587-4d73-b10b-8c9e747c68da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>RELATIVE_FILE_PATH</th>\n",
       "      <th>FILE_NAME</th>\n",
       "      <th>SAMPLE_RATE</th>\n",
       "      <th>SOURCE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>g1/JS00001_0.pkl</td>\n",
       "      <td>JS00001_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>g1/JS00002_0.pkl</td>\n",
       "      <td>JS00002_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>g1/JS00004_0.pkl</td>\n",
       "      <td>JS00004_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>g1/JS00005_0.pkl</td>\n",
       "      <td>JS00005_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>g1/JS00006_0.pkl</td>\n",
       "      <td>JS00006_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10242</th>\n",
       "      <td>g9/JS09367_0.pkl</td>\n",
       "      <td>JS09367_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10243</th>\n",
       "      <td>g9/JS09369_0.pkl</td>\n",
       "      <td>JS09369_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10244</th>\n",
       "      <td>g9/JS09370_0.pkl</td>\n",
       "      <td>JS09370_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10245</th>\n",
       "      <td>g9/JS09371_0.pkl</td>\n",
       "      <td>JS09371_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10246</th>\n",
       "      <td>g9/JS09372_0.pkl</td>\n",
       "      <td>JS09372_0.pkl</td>\n",
       "      <td>500</td>\n",
       "      <td>g9</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>10247 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      RELATIVE_FILE_PATH      FILE_NAME  SAMPLE_RATE SOURCE\n",
       "0       g1/JS00001_0.pkl  JS00001_0.pkl          500     g1\n",
       "1       g1/JS00002_0.pkl  JS00002_0.pkl          500     g1\n",
       "2       g1/JS00004_0.pkl  JS00004_0.pkl          500     g1\n",
       "3       g1/JS00005_0.pkl  JS00005_0.pkl          500     g1\n",
       "4       g1/JS00006_0.pkl  JS00006_0.pkl          500     g1\n",
       "...                  ...            ...          ...    ...\n",
       "10242   g9/JS09367_0.pkl  JS09367_0.pkl          500     g9\n",
       "10243   g9/JS09369_0.pkl  JS09369_0.pkl          500     g9\n",
       "10244   g9/JS09370_0.pkl  JS09370_0.pkl          500     g9\n",
       "10245   g9/JS09371_0.pkl  JS09371_0.pkl          500     g9\n",
       "10246   g9/JS09372_0.pkl  JS09372_0.pkl          500     g9\n",
       "\n",
       "[10247 rows x 4 columns]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "628c586f-b93e-40de-a931-61003405b633",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32d2a9b9-27e5-4454-b245-b5d9fdd6c2a2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a9fd448-58b2-48c5-86e0-c2d4f632d2db",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2b6aab8-9d33-4d4c-85dc-4123ccf7ff7a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e0d95ab-d658-47d6-a21b-8693bccfa614",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44de267-612f-49f0-a0c4-014f0a3c9066",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "145e3b32-4569-4098-b8f9-05967f08d266",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bed473c-f05b-40af-a2d5-320b009589f6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a7fe86-4137-437a-a1f5-4c6d781daef6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0751a74c-96c9-4f1c-b3d5-000636565096",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adf5d001-0cc0-4772-b232-7f84b35995ec",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d00647b-9456-4fe1-b189-a10633064f56",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f93b79-336d-43cb-9b07-9075e8c9d0c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save all the cropped signals\n",
    "num_saved = 0\n",
    "for record_rel_path in tqdm(record_rel_paths):\n",
    "    record_rel_dir, record_name = os.path.split(record_rel_path)\n",
    "    save_dir = os.path.join(args.output_dir, record_rel_dir)\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    source_name = record_rel_dir.split(\"/\")[0]\n",
    "    signal, record_info = wfdb.rdsamp(os.path.join(args.input_dir, record_rel_path))\n",
    "    lead_idx = np.array([record_info[\"sig_name\"].index(lead_name) for lead_name in _LEAD_NAMES])\n",
    "    signal = signal[:, lead_idx]\n",
    "    fs = record_info[\"fs\"]\n",
    "    signal_length = record_info[\"sig_len\"]\n",
    "    # if signal_length < 10 * fs:  # Exclude the ECGs with lengths of less than 10 seconds\n",
    "    #     continue\n",
    "    # cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "    # for idx, cropped_signal in enumerate(cropped_signals):\n",
    "    #     if cropped_signal.shape[1] != 10 * fs or np.isnan(cropped_signal).any():\n",
    "    #         continue\n",
    "    #     pd.to_pickle(cropped_signal.astype(np.float32),\n",
    "    #                  os.path.join(save_dir, f\"{record_name}_{idx}.pkl\"))\n",
    "    #     index_df.loc[num_saved] = [f\"{record_rel_path}_{idx}.pkl\",\n",
    "    #                                f\"{record_name}_{idx}.pkl\",\n",
    "    #                                fs,\n",
    "    #                                source_name]\n",
    "    #     num_saved += 1\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acc7bac7-3a8b-485f-a8ae-66f7cecb3ce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "record_rel_dir, record_name, save_dir, source_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1ee7f9b-5f69-493a-b3ed-2d11ec010d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67d5325a-e0d3-43d2-a882-9d00cbb02b47",
   "metadata": {},
   "outputs": [],
   "source": [
    "record_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff5b601-6f4d-4004-8b15-162aa5675b4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "if signal_length < 10 * fs:  # Exclude the ECGs with lengths of less than 10 seconds\n",
    "    continue\n",
    "cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "for idx, cropped_signal in enumerate(cropped_signals):\n",
    "    if cropped_signal.shape[1] != 10 * fs or np.isnan(cropped_signal).any():\n",
    "        continue\n",
    "    pd.to_pickle(cropped_signal.astype(np.float32),\n",
    "                 os.path.join(save_dir, f\"{record_name}_{idx}.pkl\"))\n",
    "    index_df.loc[num_saved] = [f\"{record_rel_path}_{idx}.pkl\",\n",
    "                               f\"{record_name}_{idx}.pkl\",\n",
    "                               fs,\n",
    "                               source_name]\n",
    "    num_saved += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1521670-34dd-4150-9d70-2d43c9369593",
   "metadata": {},
   "outputs": [],
   "source": [
    "if signal_length < 10 * fs:\n",
    "    print(1)\n",
    "else: print(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd94d4d2-a3ef-4966-9c63-0ee79b8ca72b",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5580e89a-0341-424a-95cc-639a29be1a6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fs * 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cacee1e-f867-487e-b215-be5f364aa365",
   "metadata": {},
   "outputs": [],
   "source": [
    "cropped_signals = moving_window_crop(signal.T, crop_length=10 * fs, crop_stride=10 * fs)\n",
    "cropped_signals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c86396f1-2cd3-4056-b0a0-9a1bd81042bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(signal).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "123cec3f-1d7f-4073-85e5-d03c8dfddf8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(cropped_signals).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bc5c577-af73-435f-b278-b79b700baaed",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, cropped_signal in enumerate(cropped_signals):\n",
    "    print(cropped_signal)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c15bd445-ac06-4e61-a41b-4ffc91df9395",
   "metadata": {},
   "outputs": [],
   "source": [
    "[f\"{record_rel_path}_{idx}.pkl\", f\"{record_name}_{idx}.pkl\", fs, source_name]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81139e44-8652-4d9a-83e9-0dabab4f0c6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = pd.read_csv('/tf/physionet.org/files/ptb-xl/1.0.3/ptbxl_database.csv', index_col='ecg_id')\n",
    "Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efc3e165-3730-4b01-b336-8be6a3246514",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y.scp_codes # = Y.scp_codes.apply(lambda x: ast.literal_eval(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6c68783-152a-4e85-8cc6-6288905f0b29",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ast\n",
    "Y.scp_codes = Y.scp_codes.apply(lambda x: ast.literal_eval(x))\n",
    "Y.scp_codes"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
